import matplotlib.pyplot as plt
import torch
import numpy as np

def plot_feature_maps(tensor, title="Feature Maps", figsize=(12, 8)):
    """
    可视化特征图
    Args:
        tensor: 形状为 [channels, height, width] 的张量
        title: 图表标题
        figsize: 图表大小
    """
    if isinstance(tensor, torch.Tensor):
        tensor = tensor.detach().cpu().numpy()
    
    n_channels = tensor.shape[0]
    n_cols = 4
    n_rows = (n_channels + n_cols - 1) // n_cols
    
    plt.figure(figsize=figsize)
    for i in range(n_channels):
        plt.subplot(n_rows, n_cols, i + 1)
        plt.imshow(tensor[i], cmap='viridis')
        plt.axis('off')
        plt.title(f'Channel {i}')
    
    plt.suptitle(title)
    plt.tight_layout()
    plt.show()

def plot_comparison(input_tensor, output_tensor, title="Input vs Output"):
    """
    对比输入和输出的特征图
    Args:
        input_tensor: 输入张量
        output_tensor: 输出张量
        title: 图表标题
    """
    if isinstance(input_tensor, torch.Tensor):
        input_tensor = input_tensor.detach().cpu().numpy()
    if isinstance(output_tensor, torch.Tensor):
        output_tensor = output_tensor.detach().cpu().numpy()
    
    plt.figure(figsize=(12, 5))
    
    plt.subplot(1, 2, 1)
    plt.imshow(np.mean(input_tensor, axis=0), cmap='viridis')
    plt.title('Input Feature Map (Mean)')
    plt.axis('off')
    
    plt.subplot(1, 2, 2)
    plt.imshow(np.mean(output_tensor, axis=0), cmap='viridis')
    plt.title('Output Feature Map (Mean)')
    plt.axis('off')
    
    plt.suptitle(title)
    plt.tight_layout()
    plt.show()

def plot_metrics(metrics_dict, title="Training Metrics"):
    """
    绘制训练指标
    Args:
        metrics_dict: 包含指标的字典，格式为 {metric_name: [values]}
        title: 图表标题
    """
    plt.figure(figsize=(10, 6))
    
    for metric_name, values in metrics_dict.items():
        plt.plot(values, label=metric_name)
    
    plt.title(title)
    plt.xlabel('Iterations')
    plt.ylabel('Value')
    plt.legend()
    plt.grid(True)
    plt.show() 